package at.ltasat.classifier.knn;

import java.util.Arrays;
import java.util.HashMap;
import java.util.Map;

import at.ltasat.classifier.Classifier;
import at.ltasat.classifier.ClazzManager;
import at.ltasat.classifier.Membership;
import at.ltasat.classifier.ObjectFeatures;
import at.ltasat.classifier.ObjectFeaturesFactory;
import at.ltasat.classifier.Trainer;

final class KNearestNeighbourTrainer<O, D> extends Trainer<O, D> {

	private final int k;
	private final int factor;
	private final Map<Point, Membership<D>[]> data;

	public KNearestNeighbourTrainer(String name, ClazzManager<D> classManager, ObjectFeaturesFactory<O> factory, int k, int factor) {
		super(name, classManager, factory);
		assert k >= 1;
		assert factor >= 1;
		this.k = k;
		this.factor = factor;
		data = new HashMap<Point, Membership<D>[]>();
	}

	@Override
	public void train(ObjectFeatures features, Membership<D>... memberships) {
		assert features != null;
		assert memberships != null;
		assert memberships.length > 0;
		assert Membership.check(memberships);
		memberships = memberships.clone();
		Point point = Point.create(features, factor);
		if (data.containsKey(point)) {
			Membership<D>[] oldMemberships = data.get(point);
			Membership<D>[] newMemberships = Arrays.copyOf(oldMemberships, oldMemberships.length + memberships.length);
			for (int i = oldMemberships.length; i < newMemberships.length; i++) {
				assert newMemberships[i] == null;
				newMemberships[i] = (Membership<D>) memberships[i - oldMemberships.length];
			}
			memberships = newMemberships;
			for (int i = 0; i < memberships.length; i++) {
				Membership<D> old = (Membership<D>) memberships[i];
				memberships[i] = new Membership<D>(old.getClazz(), old.getDegree() / 2);
			}
		}
		assert Membership.check(memberships);
		data.put(point, memberships);
	}

	@Override
	public Classifier<O, D> create() {
		return new KNearestNeighbourClassifier<O, D>(getName(), getClazzManager(), getObjectFeaturesFactory(), k, factor, data);
	}
}
